import os
from tpuSk import *
from infer_engine import InferEngine


class TPU:
    def __init__(self):

        self.tpu = tpuSk()
        # 获取当前脚本的目录路径
        self.pwd = os.path.dirname(os.path.abspath(__file__))

    # 识别链接TPU
    def link(self):
        try:
            res = self.tpu.generate_usb_config_ini(f"{self.pwd}/usbprop.ini")
            if res == 0:
                print("TPU链接成功!")
            else:
                print("TPU识别错误编号:", res)
                return False
            os.environ['USB_CONFIG_INI'] = os.path.join(self.pwd, './usbprop.ini')
            self.tpu.init()
        except Exception as e:
            print(f"TPU无法识别, error:{str(e)}")
        return True

    # TPU模型初始化
    def model_init(self, model_path,host=False):
        try:
            if host ==False:
                pass
            else:
                net_file = "host:" + f"{model_path}/net.bin"
                model_file = "host:" + f"{model_path}/model.bin"
                self.engine = InferEngine(net_file, model_file, max_batch=1)
        except Exception as e:
            print(f"模型初始化失败, error:{str(e)}")

    # 释放TPU
    def release(self):
        print("TPU释放成功")

    # 重启TPU
    def reboot(self):
        self.tpu.send_cmd("reboot")

    # 获取TPU内部的配置文件
    def get_config(self, app_name, dst_path):
        self.link()
        os.makedirs(dst_path, exist_ok=True)
        src_path = f"/root/data/model/{app_name}/model.conf"
        res = self.tpu.get_file(src_path, f"{dst_path}/model.conf", 5000)
        if res != 0:
            return False
        return True

    # 获取TPU内部的模型文件
    def get_model(self, app_name, dst_path):
        self.link()
        os.makedirs(dst_path, exist_ok=True)
        src_path = f"/root/data/model/{app_name}/net.bin"
        res1 = self.tpu.get_file(src_path, f"{dst_path}/net.bin", 5000)

        src_path = f"/root/data/model/{app_name}/model.bin"
        res2 = self.tpu.get_file(src_path, f"{dst_path}/model.bin", 8000)
        if res1 != 0 or res2 != 0:
            return False
        return True
